import os
from os.path import join

from PIL import Image
from tqdm import tqdm

from models.unet.unet import Unet
from utils.evaluation.miou import miou_Unet, compute_mIoU


def compute(netName, num_classes):
    gt_dir = "E:/python/data/VOC2012/VOCdevkit/VOC2012/SegmentationClass"
    pred_dir = "miou_pr_dir_" + netName
    png_name_list = open("E:/python/data/VOC2012/VOCdevkit/VOC2012/ImageSets/Segmentation/val.txt",
                         'r').read().splitlines()

    # ------------------------------------------------#
    #   Obtain a list of verification set tag paths, which is easy to read directly
    #   Obtain the path list of the image segmentation results of the validation set,
    #   which is convenient for direct reading
    # ------------------------------------------------#
    gt_imgs = [join(gt_dir, x + ".png") for x in png_name_list]
    for ind in range(len(gt_imgs)):
        gt_imgs[ind] = Image.open(gt_imgs[ind])
    pred_imgs = [join(pred_dir, x + ".png") for x in png_name_list]

    name_classes = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow",
                    "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train",
                    "tvmonitor"]
    compute_mIoU(gt_imgs, pred_imgs, num_classes, name_classes)  # 执行计算mIoU的函数


def init(net_name, num_classes, model_path, model_image_size):
    unet = miou_Unet(
        net=Unet(num_classes=num_classes, in_channels=model_image_size[-1]).eval(),
        model_path=model_path,
        model_image_size=model_image_size)
    image_ids = open("E:/python/data/VOC2012/VOCdevkit/VOC2012/ImageSets/Segmentation/val.txt", 'r').read().splitlines()
    if not os.path.exists("miou_pr_dir_" + net_name):
        os.makedirs("miou_pr_dir_" + net_name)
    for image_id in tqdm(image_ids):
        image_path = "E:/python/data/VOC2012/VOCdevkit/VOC2012/JPEGImages/" + image_id + ".jpg"
        image = Image.open(image_path)
        image = unet.detect_image(image)
        image.save("miou_pr_dir_" + net_name + "/" + image_id + ".png")


def main():
    netName = "unet"
    model_path = "E:/python/python1/RVSC2012Project/VOC/checkpoint_unet_1/Epoch40-Total_Loss0.1551-Val_Loss0.2993.pth"
    model_image_size = (256, 256, 3)
    num_classes = 21

    init(netName, model_path, model_image_size)
    compute(netName, num_classes)


if __name__ == "__main__":
    main()
